/*
 * bst/splay_tree.hh - Splay tree implementation.
 * Copyright (C) 2008  Mikhail Vorozhtsov
 * See the LICENSE section of the README file for details.
 */

/* $Id$ */

#ifndef HEADER_LIBBST_MJ_SPLAY_TREE_HH
#define HEADER_LIBBST_MJ_SPLAY_TREE_HH

#include HEADER (LIBBST_PREFIX,base.hh)

START_NAMESPACE (LIBBST_NS)

template <typename NODE, typename TREE, typename GLUE>
struct fwd_splay_tree_t : public base_tree_t<NODE, TREE, GLUE>
{
protected:
  template <typename STACK>
  void
  splay (NODE *x, NODE *root, NODE *root_parent, STACK &stack)
  {
    NODE *parent;
      
    if (x == root)
      return;
    else
      parent = stack.pop (this, x);

    while (true)
      {
        if (parent == root)
          {
            if (x == parent->get_left ())
              parent->rotate_right (this, root_parent);
            else
              parent->rotate_left (this, root_parent);

            break;
          }

        NODE *grandparent = stack.pop (this, parent);
        NODE *gg = grandparent == root ? root_parent
                   : stack.pop (this, grandparent);

        if (parent == grandparent->get_left ())
          {
            if (x == parent->get_left ())
              {
                grandparent->rotate_right (this, gg);
                parent->rotate_right (this, gg);
              }
            else
              {
                parent->rotate_left (this, grandparent);
                grandparent->rotate_right (this, gg);
              }
          }
        else
          {
            if (x == parent->get_left ())
              {
                parent->rotate_right (this, grandparent);
                grandparent->rotate_left (this, gg);
              }
            else
              {
                grandparent->rotate_left (this, gg);
                parent->rotate_left (this, gg);
              }
          }

        if (grandparent == root)
          break;

        parent = gg;
      }
  }

public:
  template <typename KEY, typename CMP, typename STACK>
  typename NODE::container_t *
  lookup_key (const KEY &key, STACK &stack) const
  {
    NODE *parent, *prev_or_next;
    bool left_p;
    typename NODE::container_t *x
      = this->template full_lookup_key<KEY, CMP, STACK>
                (key, parent, left_p, prev_or_next, stack);
    
    if (x)
      {
        ((fwd_splay_tree_t *) this)->splay (
          (NODE *) NODE::from_container (x), (NODE *) this->get_root (),
          NULL, stack);
        return x;
      }

    if (parent)
      {
        stack.pop ();
        ((fwd_splay_tree_t *) this)->splay (
          parent, (NODE *) this->get_root (), NULL, stack);
      }

    return NULL;
  }

  template <typename STACK>
  bool
  insert (typename NODE::container_t *x, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;

    if (full_lookup (x, parent, left_p, prev_or_next, stack))
      return false;

    NODE *n = (NODE *) NODE::from_container (x);

    before_insert (n, parent, left_p);
    n->insert (this, parent, left_p, prev_or_next);
    splay (n, (NODE *) this->get_root (), NULL, stack);

    return true;
  }

  template <typename KEY, typename CMP, typename STACK>
  typename NODE::container_t *
  remove_key (const KEY &key, STACK &stack)
  {
    NODE *parent, *prev_or_next;
    bool left_p;
    typename NODE::container_t *x
      = this->template full_lookup_key<KEY, CMP, STACK>
                         (key, parent, left_p, prev_or_next, stack);

    if (!x)
      return NULL;

    NODE *n = (NODE *) NODE::from_container (x);

    splay (n, (NODE *) this->get_root (), NULL, stack);

    NODE *left = (NODE *) n->get_left ();
    NODE *right = (NODE *) n->get_right ();

    if (left)
      {
        NODE *pred_parent = n;
        stack.push (n);
        NODE *pred = left->get_max (pred_parent, stack);
        splay (pred, left, n, stack);
        NODE *succ = NULL;
        n->replace_with_pred (this, NULL, true, pred, n, succ);
        after_remove (n, NULL, pred, right != NULL, succ);
      }
    else if (right)
      {
        NODE *succ_parent = n;
        stack.push (n);
        NODE *succ = right->get_min (succ_parent, stack);
        splay (succ, right, n, stack);
        n->replace_with_succ (this, NULL, true, succ, n, NULL);
        after_remove (n, NULL, NULL, true, succ);
      }
    else
      {
        n->remove (this, NULL, true, NULL);
        after_remove (n, NULL, NULL, false, NULL);
      }

    return x;
  }
};

template <typename NODE, typename TREE = basic_tree_t, typename GLUE = glue_t>
struct splay_tree_t
  : public wrap_tree_t<fwd_splay_tree_t<NODE, TREE, GLUE> >
{
};

END_NAMESPACE

#endif /* HEADER_LIBBST_MJ_SPLAY_TREE_HH */

